// Copyright 2023 Srhino Co., Ltd. All rights reserved.

package v4

import (
	"bytes"
	"context"
	"errors"
	v2 "git.ouryun.cn/lzj1/grpc-bidirectional-stream/pkg/apis/v2"
	ifacev2 "git.ouryun.cn/lzj1/grpc-bidirectional-stream/pkg/interfaces/v2"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"io"
	"log"
	"net"
	"net/http"
	"sync/atomic"
	"time"
)

var reqId int64

type grpcServer struct {
	v2.UnimplementedStreamDataServiceServer
	address    string
	streamData *StreamData
	stream     v2.StreamDataService_AgentStreamServer
	ctx        context.Context
}

var _ ifacev2.GrpcServer = (*grpcServer)(nil)

func NewGrpcServer(addr string, ctx context.Context) ifacev2.GrpcServer {
	return &grpcServer{
		address:    addr,
		streamData: NewStreamData(),
	}
}

func (s *grpcServer) AgentStream(stream v2.StreamDataService_AgentStreamServer) error {
	defer func() {
		if r := recover(); r != nil {
			log.Println(r)
		}
	}()
	log.Println("client connection is successful")

	s.stream = stream

	// 接收客户端数据
	err := s.handleRecv()
	if err != nil {
		return err
	}
	return nil
}

// 接收grpc客户端数据
func (s *grpcServer) handleRecv() error {
	for {
		recv, err := s.stream.Recv()
		log.Printf("接收到客户端数据 %v", recv)
		// 错误处理
		if err != nil {
			log.Println(err)
			if grpc.Code(err) == codes.Canceled {
				s.stream = nil
			}

			if err == io.EOF {
				s.stream = nil
			}
			log.Printf("recv error : %v", err)
		}
		// 客户端数据处理
		//_, isRequest := message.Message.(*v2.StreamMessage_Request)
		if _, ok := recv.Message.(*v2.StreamMessage_Response); ok {
			err := s.handleResponse(recv)
			if err != nil {
				log.Printf("handleResponse error : %v", err)
			}
		} else {
			s.handleRequest(recv)
		}
	}

}

// 处理grpc客户端响应数据
func (s *grpcServer) handleResponse(recv *v2.StreamMessage) error {
	log.Printf("handleResponse %v", recv)
	s.streamData.Responses <- recv
	return nil
}

// 处理grpc客户端请求数据
func (s *grpcServer) handleRequest(recv *v2.StreamMessage) {
	log.Printf("handleRequest %v", recv)
	// 协议转换
	request := recv.Message.(*v2.StreamMessage_Request).Request

	resp, err := s.httpRequest(request)

	streamResponse := &v2.StreamMessage_Response{
		Response: &v2.Response{
			Code: v2.Response_SUCCESS_CODE,
		},
	}

	if err != nil {
		streamResponse.Response.Code = v2.Response_FAIL_CODE
		streamResponse.Response.Error = err.Error()
	} else {
		streamResponse.Response.Body = resp
	}

	message := &v2.StreamMessage{
		RequestId: recv.RequestId,
		Message:   streamResponse,
	}
	if err := s.sendMessage(message); err != nil {
		log.Println(err)
	}
}

// 发送http请求
func (s *grpcServer) httpRequest(request *v2.Request) ([]byte, error) {
	req, err := http.NewRequest(request.Method, request.Url, bytes.NewReader(request.Body))
	if err != nil {
		return nil, err
	}

	for k, h := range request.Headers {
		req.Header.Add(k, h)
	}

	resp, err := http.DefaultClient.Do(req)

	if err != nil {
		return nil, err
	}

	defer resp.Body.Close()

	bs, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	return bs, nil
}

// 处理转发请求
func (s *grpcServer) forwardRequest() {
	for {
		select {
		case request := <-s.streamData.Requests:
			log.Println(request)
			if err := s.sendMessage(request); err != nil {
				log.Println(err)
			}
		}
	}

}

// 发送消息到grpc客户端
func (s *grpcServer) sendMessage(message *v2.StreamMessage) error {

	stream := s.stream

	if stream == nil {
		return errors.New("stream does not exist")
	}
	log.Printf("发送消息 %v", message)
	err := stream.Send(message)

	if err != nil {
		return err
	}

	return nil
}

// Request 接收转发请求
func (s *grpcServer) Request(message *v2.StreamMessage, requestId uint64) (*v2.StreamMessage, error) {
	_, isRequest := message.Message.(*v2.StreamMessage_Request)

	if isRequest {
		atomic.AddInt64(&reqId, 1)
		message.RequestId = uint64(reqId)
	} else {
		message.RequestId = requestId
	}

	s.streamData.Requests <- message
	if isRequest {
		for {
			select {
			case <-time.After(5 * time.Second):
				return nil, errors.New("request timeout")
			case msg := <-s.streamData.Responses:
				// TODO 这里的处理方式还需要优化
				if msg.RequestId == message.RequestId {
					return msg, nil
				}
			}
		}
	}
	return nil, nil
}

func (s *grpcServer) Start() error {
	go func() {
		s.forwardRequest()
	}()
	log.Println("grpc server start")
	listener, _ := net.Listen("tcp", s.address)
	grpcserver := grpc.NewServer(
		// 设置最大接收和发送消息大小
		grpc.MaxSendMsgSize(100<<20), // 100 MB
		grpc.MaxRecvMsgSize(100<<20), // 100 MB
	)
	// 注册服务方法
	v2.RegisterStreamDataServiceServer(grpcserver, s)

	if err := grpcserver.Serve(listener); err != nil {
		log.Println(err)
		return err
	}

	return nil
}

func (s *grpcServer) Stop() {
	// TODO: 关闭连接
}
